import os

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import  ChatOpenAI
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import OpenAIEmbeddings

# 配置http和https的服务器代理
os.environ["http_proxy"] = "127.0.0.1:8888"
os.environ["https_proxy"] ="127.0.0.1:8888"
# 配置openai的密钥
os.environ["OPENAI_API_KEY"] = "你的密钥"

# 创建deepseek模型
model = ChatOpenAI(model="deepseek-chat",api_key="你的密钥",base_url= "https://api.deepseek.com")

# 创建一个文档列表
documents = [
    Document(
        page_content="《肖申克的救赎》讲述了一位被错误判刑的银行家安迪·杜佛兰，在肖申克监狱中展开了一场漫长的自我救赎之旅，最终通过智慧和毅力重获自由的故事。",
        metadata={"title": "肖申克的救赎", "year": 1994, "director": "弗兰克·德拉邦特", "genre": "剧情"}
    ),
    Document(
        page_content="《教父》是一部经典的黑帮电影，讲述了意大利裔美国人维托·柯里昂如何建立并维持他的犯罪帝国的故事，揭示了权力、家庭和背叛的主题。",
        metadata={"title": "教父", "year": 1972, "director": "弗朗西斯·福特·科波拉", "genre": "犯罪"}
    ),
    Document(
        page_content="《阿甘正传》通过一个智商只有75的男人的视角，讲述了美国历史上从20世纪50年代到90年代的重要事件，展现了普通人的非凡人生和深刻的人生哲理。",
        metadata={"title": "阿甘正传", "year": 1994, "director": "罗伯特·泽米吉斯", "genre": "剧情"}
    ),
    Document(
        page_content="《泰坦尼克号》是一部史诗般的爱情电影，讲述了1912年泰坦尼克号沉船事件中，两位来自不同社会阶层的年轻人——杰克和露丝的感人爱情故事。",
        metadata={"title": "泰坦尼克号", "year": 1997, "director": "詹姆斯·卡梅隆", "genre": "爱情"}
    ),
    Document(
        page_content="《盗梦空间》是一部复杂的科幻电影，讲述了一群盗梦者通过潜入他人梦境进行盗窃的故事，层层递进的梦境结构和伦理道德的探讨令观众深思。",
        metadata={"title": "盗梦空间", "year": 2010, "director": "克里斯托弗·诺兰", "genre": "科幻"}
    )
]

# 实例化一个向量数据库
vector_store = Chroma.from_documents(documents, embedding = OpenAIEmbeddings())

#进行相似度查询，返回相似分数，分数越低相似度越高
# print(vector_store.similarity_search_with_score("爱情电影"))

#配置检索器
retriever = RunnableLambda(vector_store.similarity_search).bind(k=1)

# 查询出要检索的内容
# print(retriever.batch("爱情"))

# 创建一个message长字符串
message = """
根据上下文回答这个问题：
{question}
上下文：
{context}
"""

# 接着创建提示词模板
prompt_temp = ChatPromptTemplate.from_messages([("human", message)])

#创建chain，期中RunnablePassthrough允许我们将问题传给prompt和model
#这段代码定义了一个链式结构，用于根据用户的问题从上下文中检索信息，并生成回答。
# RunnablePassthrough() 用于直接传递用户输入的问题。
# retriever 用于检索与问题相关的上下文。
# prompt_temp 用于构建提示词。
# model 用于根据提示词生成回答。
chain = {'question':RunnablePassthrough(),'context':retriever}|prompt_temp|model

# 执行chain并打印出模型给的结果
print(chain.invoke("《泰坦尼克号》的介绍").content)








